#include "covariance.h"
#include <iostream>
#include <iomanip>
#include <cassert>
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <sstream>
#include "vector_tools.h"

using namespace std;

Covariance::Covariance(int rank_): rank(rank_),
  sum(0), sum2(0)
{
  //demand nonzero rank
  assert(rank_>0);

  //initialize matrix
  ary2.resize(rank*rank);
  ary.resize(rank);
}

bool Covariance::Read(char* input)
{
  //first find the string "full xy matrix"
  ifstream fin_matrix(input);
  string line;
  bool read_matrix=false;

  while(fin_matrix.good())
    {
      getline(fin_matrix, line);

      //we have a match
      if(line.find("full xy matrix")!=string::npos)
	{
	  //now read the file
	  for(int i=0; i<rank; ++i)
	    {
	      getline(fin_matrix, line);

	      //now process that particular line
	      stringstream ssout(line);
	      string token;

	      //cout<<"processing line: "<<line<<endl;

	      //get all the numbers in line
	      for(int j=0; j<rank; ++j)
		{
		  if(!getline(ssout, token, ','))
		    {
		      cout<<"error line: "<<line<<endl;
		      cout<<"ssout: "<<ssout<<endl;
		      cout<<"token: "<<token<<endl;
		      ary2[i*rank+j]=0;
		      //throw;
		    }
		    

		  //cout<<"(i, j): "<<i<<", "<<j<<", "<<token<<endl;
		  ary2[i*rank+j]=atof(token.c_str());
		}
	    }

	  read_matrix=true;
	  break;
	}
      
    }
  
  //if failed to read matrix
  /*
  if(!read_matrix)
    throw;
  */

  bool read_vec=false;
  
  //now get the ary2 vector
  while(fin_matrix.good())
    {
      getline(fin_matrix, line);
      
      //we have a match
      if(line.find("full x vector")!=string::npos)
	{
	  getline(fin_matrix, line);

	  //now process that particular line
	  stringstream ssout(line);
	  string token;
	  
	  //get all the numbers in line
	  for(int j=0; j<rank; ++j)
	    {
	      if(!getline(ssout, token, ','))
		throw;
	      
	      ary[j]=atof(token.c_str());
	    }

	  read_vec=true;
	  break;
	}
    }

  //if failed to read vector
  /*
  if(!read_vec)
    throw;
  */

  //now read sum^2
  while(fin_matrix.good())
    {
      getline(fin_matrix, line);

      //we have a match
      if(line.find("sum^2")!=string::npos)
	{

	  string subline=line.substr(6);
	  sum2=atof(subline.c_str());	
	  break;
	}

    }  

  //now read sum
  while(fin_matrix.good())
    {
      getline(fin_matrix, line);
      
      //we have a match
      if(line.find("sum")!=string::npos)
	{

	  string subline=line.substr(4);
	  sum=atof(subline.c_str());	  
	  break;
	}

    }  


  fin_matrix.close();
  return true;
}

Covariance& Covariance::operator+=(const Covariance& input)
{
  //make sure rank is the same
  assert(rank==input.rank);

  sum+= input.sum;
  sum2+= input.sum2;
  
  for(int i=0; i<rank; ++i)
    {
      ary[i]+=input.ary[i];
      
      for(int j=0; j<rank; ++j)
	{
	  ary2[i*rank+j]+=input.ary2[i*rank+j];
	}
    }
  return *this;
}

double Covariance::mean(int row) const
{
  assert(row>=0 && row<rank);

  //make sure entry isn't zero
  assert(sum2>0);

  double meanx=ary[row]/sum;

  return meanx;
}

double Covariance::operator()(int row, int col) const
{
  //demand indices to be within range
  assert(row>=0 && row<rank);
  assert(col>=0 && col<rank);
  //make sure entry isn't zero
  assert(sum2>0);

  //effective number of entries
  double eff_n=sum*sum/sum2;

  //apply N/(N-1) to get an unbiased estimate
  //if eff_n < 2, just ignore such factor
  double unbias_factor=1.0;
  if(eff_n>2) unbias_factor=eff_n/(eff_n-1);

  double meanxy=ary2[row*rank+col]/sum;
  double meanx=ary[row]/sum;
  double meany=ary[col]/sum; 
  
  return unbias_factor*(meanxy-meanx*meany);
}

void Covariance::Fill(const vector<double>& ntuple, double weight)
{
  //make sure ntuple is of right size
  assert(ntuple.size()==rank);
  //increment counts
  sum2+=weight*weight;
  sum+=weight;
  
  //loop over ary2 and ary and increment x*y and x
  for(int i=0; i<rank; i++)
    {
      ary[i]+=(ntuple[i]*weight);
      for(int j=i; j<rank; j++)
	{
	  ary2[i*rank+j]+= (ntuple[i]*ntuple[j]*weight);
	  //symmetrize results
	  ary2[j*rank+i]=ary2[i*rank+j];
	}
    }
}


ostream& operator<<(ostream& os, const Covariance& co)
{
  //cout covariance matrix

  /*
  os<<"covariance: "<<endl;
  os<<endl;
  for(int i=0; i<co.rank; i++)
    {
	for(int j=0; j<co.rank; j++)
	  {
	    os<<setw(5)<<setprecision(10)<<co(i,j);
	    if(j<co.rank-1)
	      cout<<", ";
	  }
	cout<<endl;
    }

  os<<endl;
  */

  os<<"full xy matrix: "<<endl;
  for(int i=0; i<co.rank; i++)
    {
	for(int j=0; j<co.rank; j++)
	  {
	    //os<<setw(5)<<setprecision(3)
	    os<<co.ary2[i*co.rank+j];
	    if(j<co.rank-1)
	      cout<<", ";
	  }
	cout<<endl;
    }
  os<<endl;

  os<<"full x vector: "<<endl;
  for(int i=0; i<co.rank; i++)
    {
      //os<<setw(5)<<setprecision(3)
      os<<co.ary[i];
      if(i<co.rank-1)
	cout<<", ";
    }
  os<<endl;
 
  os<<"sum^2: "<<co.sum2<<endl;

  os<<"sum: "<<co.sum<<endl;

  //cout bandwidth * covariance
  //to check and see if it is proportional to unity
  /*
  for(int i=0; i<co.rank; i++)
    {
      os<<"[ ";
	for(int j=0; j<co.rank; j++)
	  {
	    double value=0;
	    for(int r=0; r<co.rank; r++)
	      {
		value+=co(i,r)*co.bandwidth[r*co.rank+j];
	      }
	    os<<setw(5)<<setprecision(3)<<value<<" ";

	  }
      os<<"]"<<endl;
    }
  */
  return os;
}
